from .common import *

# draw a breakdown chart.
def DrawBreakdownFig(x_values, y_values, legend_labels, x_label, y_label,
                     x_min, x_max, y_min, y_max, bar_width=0.4,
                     folder='./', filename='default', allow_legend=True,
                     color_map=COLOR_MAP, patterns=PATTERNS):
    """DrawBreakdownFig: Function used to draw breakdown figures.
    \t Usage:
    \t\t x_values: List of x tick values.
    \t\t y_values: values of each breakdown segment.
    \t\t legend_labels: Legend labels.
    \t\t x_label: Label for the x axis.
    \t\t y_label: Label for the y axis.
    \t\t x_{min, max}: The min and max value along the x axis.
    \t\t y_{min, max}: The min and max value along the y axis.
    \t\t bar_width: The width of each bar.
    \t\t folder: The folder name.
    \t\t filename: The filename.
    \t\t allow_legend: The legend flag.
    \t\t color_map: The color map.
    \t\t patterns: The pattern set.
    """
    assert(len(legend_labels) >= len(y_values))

    # you may change the figure size on your own
    fig = plt.figure(figsize=(8, 3.2))
    figure = fig.add_subplot(111)

    # values in the x_axis
    index = np.arange(1, len(x_values) + 1)
    # draw the bars
    bars = [None] * (len(legend_labels))

    bottom_base = np.zeros(len(y_values[0]))
    for i in range(len(y_values)):
        bars[i] = figure.bar(index, y_values[i], bar_width, hatch=patterns[i],
                             color=color_map[i], label=legend_labels[i],
                             edgecolor='black', bottom=bottom_base)
        bottom_base = np.array(y_values[i]) + bottom_base

    # sometimes you may not want to draw legends.
    if allow_legend:
        figure.legend(bars, legend_labels, prop=LEGEND_FP, loc='upper center',
                      ncol=len(y_values), bbox_to_anchor=(0.45, 1.2),
                      shadow=False, frameon=False, borderaxespad=-0.3,
                      handlelength=2, labelspacing=0.2)

    # you may need to tune the xticks position to get the best figure.
    plt.xticks(index, x_values)
    y_ticks = np.arange(y_min, y_max, 0.1)
    if isinstance(y_max, int):
        if y_max < 3:
            y_max = math.ceil(y_max / 0.5) * 0.5
            y_ticks = np.arange(y_min, y_max + 0.1, 0.5)
        elif y_max >= 3 and y_max < 9:
            y_max = math.ceil(y_max / 1) * 1
            y_ticks = np.arange(y_min, y_max + 1, 1)
        else:
            y_max = math.ceil(y_max / 2) * 2
            y_ticks = np.arange(y_min, y_max + 1, 2)
    elif isinstance(y_max, float):
        if y_max < 3.0:
            y_max = math.ceil(y_max / 0.5) * 0.5
            y_ticks = np.arange(y_min, y_max + 0.1, 0.5)
        elif y_max < 8.0:
            y_max = math.ceil(y_max / 1.0) * 1.0
            y_ticks = np.arange(y_min, y_max + 0.1, 1.0)
        else:
            y_max = math.ceil(y_max / 2.0) * 2.0
            y_ticks = np.arange(y_min, y_max + 0.1, 2.0)
    else:
        print("UNKNOWN TYPES!")
        exit()
    plt.yticks(y_ticks, y_ticks)

    plt.xlim(x_min, x_max)
    plt.ylim(y_min, y_max)

    figure.grid(axis='y', color='gray', linestyle='--')
    figure.yaxis.set_major_locator(LinearLocator(len(y_ticks)))

    figure.get_xaxis().set_tick_params(direction='in', pad=10)
    figure.get_yaxis().set_tick_params(direction='in', pad=10)

    plt.xlabel(x_label, fontproperties=LABEL_FP)
    plt.ylabel(y_label, fontproperties=LABEL_FP)

    if folder and filename:
        if not os.path.exists(folder):
            os.makedirs(folder)
        plt.savefig(folder + '/' + filename + '.eps', bbox_inches='tight',
                    format='eps')
